"""
@Desc: 趋势跟踪策略的逻辑
趋势识别：
    使用技术指标（如均线、MACD、布林带等）识别市场趋势。
    例如：当价格突破某一均线时，认为趋势形成。
入场信号：
    当趋势向上时，买入。
    当趋势向下时，卖出或做空。
出场信号：
    当趋势反转或达到止损/止盈条件时，平仓。
风险管理：
    设置止损和止盈，控制单笔交易的风险。

@Auth: meihongliang-m2
@Date: 2025/3/18-16:01
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import yfinance as yf  # 用于获取股票数据


# 1. 获取数据
def get_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date)
    return data['Adj Close']  # 使用调整后的收盘价


# 2. 计算均线
def calculate_moving_averages(data, short_window, long_window):
    data['Short_MA'] = data['Adj Close'].rolling(window=short_window, min_periods=1).mean()
    data['Long_MA'] = data['Adj Close'].rolling(window=long_window, min_periods=1).mean()
    return data


# 3. 生成交易信号
def generate_signals(data):
    data['Signal'] = 0  # 初始化信号列
    data['Signal'][short_window:] = np.where(data['Short_MA'][short_window:] > data['Long_MA'][short_window:], 1, -1)
    data['Position'] = data['Signal'].diff()  # 计算仓位变化（买入或卖出）
    return data


# 4. 回测策略
def backtest(data):
    initial_capital = 10000  # 初始资金
    position = 0  # 持仓数量
    portfolio_value = []  # 记录每日投资组合价值

    for i in range(len(data)):
        if data['Position'][i] == 1:  # 买入信号
            position = initial_capital // data['Adj Close'][i]  # 计算买入数量
            initial_capital -= position * data['Adj Close'][i]  # 更新剩余资金
        elif data['Position'][i] == -1:  # 卖出信号
            initial_capital += position * data['Adj Close'][i]  # 更新资金
            position = 0  # 清空持仓

        portfolio_value.append(initial_capital + position * data['Adj Close'][i])  # 计算每日组合价值

    data['Portfolio_Value'] = portfolio_value
    return data


# 5. 可视化结果
def plot_results(data):
    plt.figure(figsize=(14, 7))
    plt.plot(data['Adj Close'], label='Price', alpha=0.5)
    plt.plot(data['Short_MA'], label=f'Short MA ({short_window})', alpha=0.75)
    plt.plot(data['Long_MA'], label=f'Long MA ({long_window})', alpha=0.75)
    plt.scatter(data.index, data['Adj Close'], c=data['Position'], cmap='viridis', label='Buy/Sell Signal', marker='^',
                s=100)
    plt.title('Trend Following Strategy (Moving Average Crossover)')
    plt.legend()
    plt.show()


# 主程序
if __name__ == "__main__":
    # 参数设置
    ticker = 'AAPL'  # 股票代码
    start_date = '2020-01-01'
    end_date = '2023-01-01'
    short_window = 10  # 短期均线窗口
    long_window = 50  # 长期均线窗口

    # 获取数据
    data = get_data(ticker, start_date, end_date)
    data = pd.DataFrame(data)
    # 计算均线
    data = calculate_moving_averages(data, short_window, long_window)
    # 生成信号
    data = generate_signals(data)
    # 回测策略
    data = backtest(data)
    # 可视化结果
    plot_results(data)
    # 打印最终收益
    final_portfolio_value = data['Portfolio_Value'].iloc[-1]
    print(f"Final Portfolio Value: {final_portfolio_value:.2f}")
